# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import recordclass
import functools
import os
import datetime
import logging
import sys
import numpy
import pandas

from absl import app, flags

from concurrent.futures import ProcessPoolExecutor
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.huobi_futures.kr_rest.futures_product import HuobiFuturesProduct
from coin.support.proto_log.logic.util import run_from_og_log_archive
from coin.strategy.mm.dumper_base import run_from_archive
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str
from coin.strategy.mm.tool.order_reader import OrderReader, OrderInfo
from coin.strategy.mm.fastfeature.ridge_regress import get_quantile

import coin.strategy.mm.feed as feed
import coin.strategy.mm.subscription as sub

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

FLAGS = flags.FLAGS

OrderIter = recordclass.recordclass('OrderIter', ['tag', 'orders', 'idx'])


def statstr(statrow):
  return "%5.2f(%4d/%4d)" % (statrow.mean(), statrow.sum(), len(statrow))


def autoquantile(xs, ys, max_quantile):
  sel = ~numpy.isnan(xs) & ~numpy.isnan(ys)
  xs = numpy.array(xs)[sel]
  ys = numpy.array(ys)[sel]
  unique_qs = len(numpy.unique(xs))
  if unique_qs < max_quantile * 2:
    df = pandas.DataFrame({'x': xs, 'y': ys})
    meandf = df.groupby('x').mean()
    xq = numpy.array(meandf.index)
    yq = numpy.array(meandf)
    return xq, yq
  else:
    return get_quantile(xs, ys, max_quantile)


def get_bid_ask_qty_by_bp(book, ret_bp):
  midp = 0.5 * (book.ask0().price + book.bid0().price)

  qtys = []
  for lb, ub, iterator in [(midp, midp * (1 + ret_bp * 1e-4), book.iteritems_ask()),
                           (midp * (1 - ret_bp * 1e-4), midp, book.iteritems_bid())]:
    qty_sum = 0
    for price, level in iterator:
      if lb <= price <= ub:
        qty_sum += level.qty
      else:
        break
    qtys.append(qty_sum)

  return {'ask%sbp_qty' % ret_bp: qtys[0], 'bid%sbp_qty' % ret_bp: qtys[1]}


def get_liquidity_by_price(book, sign, order_price):
  qty = 0
  if sign == 1:
    for price, level in book.iteritems_ask():
      if price > order_price:
        break
      qty += level.qty
  else:
    for price, level in book.iteritems_bid():
      if price < order_price:
        break
      qty += level.qty
  return qty


def plot_print_df(df, symbol, tdstr):
  plt.rcParams['font.size'] = 6
  plt.rcParams['figure.figsize'] = 10, 5
  plt.rcParams['lines.markersize'] = 1

  signdict = {-1: 'SELL', 1: 'BUY'}
  payup = df['payup']
  fill_ret = df['fill_ret']
  ret = df['ret']
  agg_select = df['payup'] >= -1e-8
  acc_agg_select = df['acc_payup'] >= -1e-8
  oppliq = df['oppliq']

  acclat = df['accepted_timestamp'] - df['timestamp']
  atflat = df['filled_timestamp'] - df['accepted_timestamp']
  taken = (((df['filled_timestamp'] - df['accepted_timestamp']) < 3e8)
           | ~numpy.isclose(df['fill_price'], df['price']))
  filled = (numpy.isnan(df['filled_timestamp']) == False)
  filltime = df['filled_timestamp'] - df['timestamp']

  fillrat = df['fill_qty'] / df['order_qty']
  aggret = df['aggret']
  mtm_price = df.loc[0, 'price']

  print(("STAT %s %s TR%s FR%s " + "aggret%6.2fbp prcret%6.2fbp takret%6.2fbp makret%6.2fbp "
         + "fillret %6.2fbp missret %6.2fbp acclat-med %5.2fms") %
        (tdstr,
         FLAGS.symbol,
         statstr(taken.loc[agg_select]),
         statstr(filled.loc[agg_select]), (aggret).mean(), (ret).mean(), (fill_ret[taken]).mean(),
         (fill_ret[~taken & filled]).mean(), (fill_ret[filled]).mean(), (ret[~filled]).mean(),
         acclat.median() * 1e-6))

  ftq, retq = autoquantile(filltime[(payup >= 0) & filled], fill_ret[(payup >= 0) & filled], 8)
  atq, aretq = autoquantile(acclat[(payup >= 0) & filled], fill_ret[(payup >= 0) & filled], 8)
  atfq, afretq = autoquantile(atflat[(payup >= 0) & filled], fill_ret[(payup >= 0) & filled], 8)

  ftq = numpy.array(ftq) * 1e-9
  atq = numpy.array(atq) * 1e-9
  atfq = numpy.array(atfq) * 1e-9
  plt.subplot(131)
  plt.plot(ftq, retq, 'g.-', linewidth=0.5)
  plt.xlabel("FILL_LATENCY")
  plt.ylabel("RET")
  plt.subplot(132)
  plt.plot(atq, aretq, 'g.-', linewidth=0.5)
  plt.xlabel("ACC_LATENCY")
  plt.ylabel("RET")
  plt.subplot(133)
  plt.plot(atfq, afretq, 'g.-', linewidth=0.5)
  plt.xlabel("ATF_LATENCY")
  plt.ylabel("RET")
  plt.tight_layout()
  plt.savefig("%s/%s_%s_lat.png" % (flags.FLAGS.out_dir, symbol, tdstr), dpi=200)
  plt.close()

  i = 0
  pq, tq = autoquantile(payup[payup >= 0], taken[payup >= 0], 4)
  pq /= mtm_price * 1e-4
  ncolumns = 4
  plt.subplot(2, 3, 1)
  plt.plot(pq, tq, 'g.-', linewidth=0.5)
  plt.xlabel("PAYUP")
  plt.ylabel("TR %s %s" % (tdstr, FLAGS.symbol))

  pq, fq = autoquantile(payup[payup >= 0], fillrat[payup >= 0], 4)
  pq /= mtm_price * 1e-4
  plt.subplot(2, 3, 4)
  plt.plot(pq, fq, 'g.-', linewidth=0.5)
  plt.xlabel("PAYUP")
  plt.ylabel("FR")

  pq, retq = autoquantile(payup[(payup >= 0) & filled], fill_ret[(payup >= 0) & filled], 4)
  pq /= mtm_price * 1e-4
  plt.twinx()
  plt.plot(pq, retq, 'r.-', linewidth=0.5)
  plt.xlabel("PAYUP")
  plt.ylabel("FILLRET(BLUE)")

  if False:
    ftq, retq = autoquantile(filltime[(payup >= 0) & taken], fill_ret[(payup >= 0) & taken], 8)
    ftq = numpy.array(ftq) * 1e-9
    plt.subplot(1, ncolumns, (ncolumns * i) + 4)
    plt.plot(ftq, retq, 'g.-', linewidth=0.5)
    plt.xlabel("(TAK)FILL_LATENCY")
    plt.ylabel("RET")

    ftq, retq = autoquantile(filltime[(payup >= 0) & ~taken & filled],
                             fill_ret[(payup >= 0) & ~taken & filled], 8)
    ftq = numpy.array(ftq) * 1e-9
    plt.subplot(1, ncolumns, (ncolumns * i) + 5)
    plt.plot(ftq, retq, 'g.-', linewidth=0.5)
    plt.xlabel("(MAK)FILL_LATENCY")
    plt.ylabel("RET")

  liqrat = numpy.clip(df['liq5bp'] / df['order_qty'], None, 20)

  pq, fq = autoquantile(liqrat, fillrat, 5)
  plt.subplot(2, 3, 2)
  plt.plot(pq, fq, 'g.-', linewidth=0.5)
  plt.xlabel("LIQ/ORDQ")
  plt.ylabel("FR - corr %s" % numpy.corrcoef(pq, fq)[0, 1])

  pq, fq = autoquantile(liqrat, fill_ret, 5)
  plt.twinx()
  plt.plot(pq, fq, 'r.-', linewidth=0.5)
  plt.ylabel("FILLRET(BLUE) - corr %s" % numpy.corrcoef(pq, fq)[0, 1])

  pq, fq = autoquantile(numpy.clip(oppliq / df['order_qty'], None, 20), fillrat, 5)
  plt.subplot(2, 3, 5)
  plt.plot(pq, fq, 'g.-', linewidth=0.5)
  plt.xlabel("OPPLIQ/ORDQ")
  plt.ylabel("FR - corr %s" % numpy.corrcoef(pq, fq)[0, 1])

  pq, fq = autoquantile(numpy.clip(oppliq / df['order_qty'], None, 20), fill_ret, 5)
  plt.twinx()
  plt.plot(pq, fq, 'r.-', linewidth=0.5)
  plt.ylabel("FILLRET(BLUE) - corr %s" % numpy.corrcoef(pq, fq)[0, 1])

  pq, fq = autoquantile(df['liq5bp'], fillrat, 5)
  plt.subplot(2, 3, 3)
  plt.plot(pq, fq, 'g.-', linewidth=0.5)
  plt.xlabel("LIQ 5bp")
  plt.ylabel("FR - corr %s" % numpy.corrcoef(pq, fq)[0, 1])

  pq, fq = autoquantile(df['liq5bp'], fill_ret, 5)
  plt.twinx()
  plt.plot(pq, fq, 'r.-', linewidth=0.5)
  plt.ylabel("FILLRET(BLUE) - corr %s" % numpy.corrcoef(pq, fq)[0, 1])

  pq, fq = autoquantile(oppliq, fillrat, 5)
  plt.subplot(2, 3, 6)
  plt.plot(pq, fq, 'g.-', linewidth=0.5)
  plt.xlabel("OPPLIQ 5bp")
  plt.ylabel("FR - corr %s" % numpy.corrcoef(pq, fq)[0, 1])

  pq, fq = autoquantile(oppliq, fill_ret, 5)
  plt.twinx()
  plt.plot(pq, fq, 'r.-', linewidth=0.5)
  plt.ylabel("FILLRET(BLUE) - corr %s" % numpy.corrcoef(pq, fq)[0, 1])

  plt.tight_layout()
  plt.savefig("%s/%s_%s_fr.png" % (flags.FLAGS.out_dir, symbol, tdstr), dpi=200)
  plt.close()


def plot_csv_stat(symbol):
  import glob
  dfs = []
  for filename in glob.glob("%s/%s_*.csv" % (flags.FLAGS.out_dir, symbol)):
    dfs.append(pandas.read_csv(filename))

  df = pandas.concat(dfs, axis=0).reset_index(drop=True)
  plot_print_df(df, symbol, "all")


class TakerRatioStrategy(object):
  def __init__(self, trading_date, product):
    self._trading_date = trading_date
    self._product = product
    # id -> (price, side, time)
    self._book_snaps = {}
    self._last_book_timestamp = 0
    self._order_iter_set = None
    self._price_info = {}
    self._order_reader = OrderReader(self._product)

  def analyze(self):
    orders_map = self._order_reader.get_orders_map()
    rescol = []
    df1 = pandas.DataFrame(list(orders_map.values()), columns=OrderInfo._fields)
    df2 = pandas.DataFrame(self._price_info).T
    df2['order_id'] = df2.index
    df = df1.merge(df2, on='order_id', how='left')

    if len(df) == 0:
      return

    if 'bid_price' not in df:
      return

    df = df.loc[~numpy.isnan(df['bid_price']) & ~numpy.isnan(df['ask_price'])
                & ~numpy.isnan(df['price'])].reset_index(drop=True)
    if len(df) == 0:
      return

    mtm_price = df.iloc[-1]['price']

    df['fill_price'] = df['fill_pq'] / df['fill_qty']

    for sign in [-1, 1]:
      midp_30s = 0.5 * (df['bid_price_30s'] + df['ask_price_30s'])
      selected = df['sign'] == sign

      if sign == -1:
        payup = df.loc[selected, 'bid_price'] - df.loc[selected, 'price']
        acc_payup = df.loc[selected, 'bid_price_accept'] - df.loc[selected, 'price']
        aggret = df.loc[selected, 'bid_price'] - midp_30s[selected]
        ret = df.loc[selected, 'price'] - midp_30s[selected]
        fill_ret = df.loc[selected, 'fill_price'] - midp_30s[selected]
        liquidity = df.loc[selected, 'liquidity']
        liq5bp = df.loc[selected, 'bid5bp_qty']
        oppliq = df.loc[selected, 'ask5bp_qty']
      elif sign == 1:
        payup = df.loc[selected, 'price'] - df.loc[selected, 'ask_price']
        acc_payup = df.loc[selected, 'price'] - df.loc[selected, 'ask_price_accept']
        aggret = midp_30s[selected] - df.loc[selected, 'ask_price']
        ret = midp_30s[selected] - df.loc[selected, 'price']
        fill_ret = midp_30s[selected] - df.loc[selected, 'fill_price']
        liquidity = df.loc[selected, 'liquidity']
        liq5bp = df.loc[selected, 'ask5bp_qty']
        oppliq = df.loc[selected, 'bid5bp_qty']
      else:
        raise ValueError(sign)

      ret = ret / midp_30s[selected] * 1e4
      fill_ret = fill_ret / midp_30s[selected] * 1e4
      aggret = aggret / midp_30s[selected] * 1e4

      payup = numpy.round(payup, 8)
      df.loc[selected, 'payup'] = payup
      df.loc[selected, 'acc_payup'] = acc_payup
      df.loc[selected, 'aggret'] = aggret
      df.loc[selected, 'ret'] = ret
      df.loc[selected, 'fill_ret'] = fill_ret
      df.loc[selected, 'liquidity'] = liquidity
      df.loc[selected, 'liq5bp'] = liq5bp
      df.loc[selected, 'oppliq'] = oppliq

    df.to_csv("oe_stat_fr/%s_%s.csv" % (FLAGS.symbol, self._trading_date.strftime("%Y%m%d")))

    plot_print_df(df, FLAGS.symbol, self._trading_date.strftime("%Y%m%d"))

  def populate_orders(self, tag):
    orders_items = []
    for order_item in list(self._order_reader._orders_map.values()):
      oitem = order_item
      if oitem.timestamp is None:
        continue
      if tag == '':
        timestamp = oitem.timestamp
      elif tag == '_accept':
        timestamp = oitem.accepted_timestamp
      elif tag == '_30s':
        timestamp = oitem.timestamp + 30 * 1e9
      else:
        raise ValueError(tag)
      if timestamp is not None:
        orders_items.append((timestamp, tag, oitem))
    return sorted(orders_items)

  def on_book_reset(self, book_builder_name, book_builder):
    product = self._product
    subscription_key = sub.get_subscription_key(product)
    if book_builder_name == subscription_key:
      book_builder.subscribe(product, functools.partial(self.on_book, product))

  def on_book(self, product, book):
    if product != self._product:
      return
    if self._last_book_timestamp + 5e8 > book.timestamp:
      return
    self._last_book_timestamp = book.timestamp

    if self._order_iter_set is None:
      self._order_iter_set = {}
      for tag in ['', '_accept', '_30s']:
        self._order_iter_set[tag] = OrderIter(tag, self.populate_orders(tag), 0)

    for tag in self._order_iter_set:
      order_iter = self._order_iter_set[tag]
      while order_iter.idx < len(order_iter.orders):
        timestamp, _, oitem = order_iter.orders[order_iter.idx]
        if timestamp < book.timestamp:
          order_iter.idx += 1
        else:
          if oitem.order_id not in self._price_info:
            self._price_info[oitem.order_id] = {}
          self._price_info[oitem.order_id].update({
              'bid_price%s' % tag: book.bid0().price, 'ask_price%s' % tag: book.ask0().price
          })
          if tag == '':
            self._price_info[oitem.order_id].update(get_bid_ask_qty_by_bp(book, 5))
            self._price_info[oitem.order_id].update(
                {'liquidity': get_liquidity_by_price(book, oitem.sign, oitem.price)})
          break

  def on_log2(self, *args, **kwargs):
    self._order_reader.on_log2(*args, **kwargs)


def launch(exchange, trading_date, root_dir, machine, symbol):
  add_hour, hours = flags.FLAGS.time_range.split("-")
  add_hour = float(add_hour)
  hours = float(hours) - add_hour
  start_time = trading_date + datetime.timedelta(hours=add_hour)

  if FLAGS.market_type == 'Spot':
    product = generate_product_from_str(FLAGS.exchange, None, symbol)
  elif FLAGS.market_type == 'Futures':
    if FLAGS.exchange == 'Okex':
      product = OkexFuturesProduct.FromStr(symbol, start_time)
    elif FLAGS.exchange == 'Huobi':
      product = HuobiFuturesProduct.FromStr(symbol, start_time)
    else:
      product = generate_product_from_str(FLAGS.exchange, None, symbol)
  else:
    raise ValueError(FLAGS.market_type)

  logging.info(trading_date)
  strategy = TakerRatioStrategy(trading_date, product)
  run_from_og_log_archive(strategy.on_log2,
                          FLAGS.market_type,
                          exchange,
                          start_time,
                          start_time + datetime.timedelta(hours=hours),
                          root_dir,
                          machine,
                          owner=None)

  if len(strategy._order_reader._orders_map) == 0:
    logging.info("%s: no order event, skip.", trading_date)
    return

  run_from_archive([product],
                   strategy.on_book_reset,
                   start_time,
                   start_time + datetime.timedelta(hours=hours),
                   machine=['feed-01.ap-northeast-1.aws'],
                   worker_id='1',
                   root_dir=FLAGS.feed_root)
  strategy.analyze()
  return 0


def get_trading_dates(trading_date_str):
  if trading_date_str.find("-") >= 0:
    tdfrom, tdto = trading_date_str.split("-")
    return pandas.date_range(tdfrom, tdto).to_pydatetime()
  else:
    return [datetime.datetime.strptime(trading_date_str, '%Y%m%d')]


def main(_):
  # import warnings; warnings.filterwarnings("error")
  os.makedirs(flags.FLAGS.out_dir, exist_ok=True)
  exchange = FLAGS.exchange
  trading_date = FLAGS.trading_date
  root_dir = os.path.expanduser(FLAGS.root_dir)
  machine = FLAGS.machine
  assert FLAGS.market_type in ('ExWallet', 'Spot', 'Margin', 'Futures')
  assert trading_date, '--trading_date must be specified.'
  assert machine, '--machine must be specified.'
  assert exchange, '--exchange must be specified.'
  print('Running for %s %s %s ...' % (trading_date, machine, exchange))
  sys.stdout.flush()
  trading_dates = get_trading_dates(trading_date)
  if len(trading_dates) == 1:
    launch(exchange, trading_dates[0], root_dir, machine, FLAGS.symbol)
  else:
    with ProcessPoolExecutor(max_workers=FLAGS.cpu) as executor:
      futures = []
      for trading_date in trading_dates:
        futures.append(
            executor.submit(launch, exchange, trading_date, root_dir, machine, FLAGS.symbol))
      for future in futures:
        future.result()
  plot_csv_stat(FLAGS.symbol)


if __name__ == '__main__':
  flags.DEFINE_string('exchange', 'Okex', 'Exchange name.')

  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')

  flags.DEFINE_string('root_dir', '~/data/og_proto_log', 'root_dir.')

  flags.DEFINE_string('feed_root', None, '')

  flags.DEFINE_string('machine', 'strategy-04.aliyun-cn-hongkong', 'Instance machine name.')

  flags.DEFINE_string('market_type', 'Futures', 'Spot, Futures, Margin')

  flags.DEFINE_string('symbol', 'EOS-USD.QUARTER', 'symbol name')

  flags.DEFINE_string('time_range', '0-24', '')

  flags.DEFINE_integer('cpu', 4, 'num cpu to use')

  flags.DEFINE_string('out_dir', 'oe_stat_fr', 'out_directory')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
